# -*- coding: utf-8 -*-
# @Function: 数据库管理器
# @Description: 负责数据库连接、表操作和数据查询，提供数据库层面的CRUD操作
# @Usage: 被importer.py、query_tester.py和llm_query_tester.py导入使用，处理数据库操作
# @Dependencies: Z_config.py

import psycopg2
from psycopg2.extras import Json
from typing import List, Tuple
from Z_config import DB_CONFIG

class DBManager:
    def __init__(self, db_config=None):
        """初始化数据库管理器

        Args:
            db_config (dict, optional): 数据库配置. Defaults to None.
        """
        self.db_config = db_config or DB_CONFIG
        self.conn = None

    def connect(self):
        """连接数据库"""
        try:
            if self.conn is None or self.conn.closed != 0:
                self.conn = psycopg2.connect(**self.db_config)
                print("数据库连接成功")
        except Exception as e:
            print(f"数据库连接失败: {e}")
            raise

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
            print("数据库连接已关闭")

    def create_table(self):
        """创建QA知识库表和表格存储表"""
        # QA知识库表
        create_qa_table_sql = """
        CREATE TABLE IF NOT EXISTS qa_knowledge (
            question TEXT PRIMARY KEY,
            question_tokens JSONB NOT NULL,
            answer JSONB NOT NULL,
            category_id VARCHAR(20) NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """

        # 表格存储表
        create_table_storage_sql = """
        CREATE TABLE IF NOT EXISTS table_storage (
            id SERIAL PRIMARY KEY,
            keywords TEXT NOT NULL,
            table_data JSONB NOT NULL,
            category_id VARCHAR(20) NOT NULL,
            source_file VARCHAR(255),
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """

        # 创建索引以提高查询性能
        create_index_sql = """
        CREATE INDEX IF NOT EXISTS idx_qa_question_tokens_gin 
        ON qa_knowledge USING GIN (question_tokens);

        CREATE INDEX IF NOT EXISTS idx_qa_tokens_array 
        ON qa_knowledge USING GIN ((question_tokens->'tokens'));

        CREATE INDEX IF NOT EXISTS idx_qa_category_id 
        ON qa_knowledge (category_id);
        
        CREATE INDEX IF NOT EXISTS idx_table_keywords 
        ON table_storage (keywords);
        
        CREATE INDEX IF NOT EXISTS idx_table_category_id 
        ON table_storage (category_id);
        """

        try:
            with self.conn.cursor() as cursor:
                cursor.execute(create_qa_table_sql)
                cursor.execute(create_table_storage_sql)
                cursor.execute(create_index_sql)
                self.conn.commit()
                print("QA知识库表和表格存储表检查/创建成功")
                return True
        except Exception as e:
            print(f"创建表失败: {e}")
            raise

    def insert_qa_data(self, qa_pairs: List[Tuple[str, str, dict]], category_id: str):
        """插入QA数据到数据库

        Args:
            qa_pairs (List[Tuple[str, str, dict]]): 问答对列表，每项包含(问题, 答案, 分词结果)
            category_id (str): 问题所属的类别ID
        """
        insert_sql = """
        INSERT INTO qa_knowledge (question, question_tokens, answer, category_id)
        VALUES (%s, %s, %s, %s)
        ON CONFLICT (question) DO UPDATE SET
            question_tokens = EXCLUDED.question_tokens,
            answer = EXCLUDED.answer,
            category_id = EXCLUDED.category_id,
            updated_at = CURRENT_TIMESTAMP;
        """

        try:
            with self.conn.cursor() as cursor:
                for question, answer, token_data in qa_pairs:
                    # 构建答案JSON
                    answer_json = {"answer": answer}

                    # 插入数据
                    cursor.execute(insert_sql, (
                        question,
                        Json(token_data),
                        Json(answer_json),
                        category_id
                    ))

                self.conn.commit()
                print(f"成功插入或更新 {len(qa_pairs)} 条QA数据")
                return True
        except Exception as e:
            print(f"插入数据失败: {e}")
            self.conn.rollback()
            raise

    def query_qa(self, tokens: List[str], limit: int = 10, category_id: str = None):
        """根据分词结果查询QA数据
    
        Args:
            tokens (List[str]): 分词结果
            limit (int, optional): 返回结果数量限制. Defaults to 5.
            category_id (str, optional): 水库类别ID. Defaults to None.
    
        Returns:
            list: 查询结果
        """
        # 基础查询SQL
        query_sql = """
        SELECT question, question_tokens, answer, category_id,
               (
                   SELECT COUNT(*)
                   FROM jsonb_array_elements_text(question_tokens->'tokens') AS token
                   WHERE token = ANY(%s)
               ) as match_count
        FROM qa_knowledge
        WHERE question_tokens->'tokens' ?| %s
        """
        
        # 如果指定了category_id，添加过滤条件
        params = [tokens, tokens]
        if category_id:
            query_sql += " AND category_id = %s"
            params.append(category_id)
            print(f"按category_id过滤: {category_id}")
        
        # 添加排序和限制
        query_sql += """
        ORDER BY match_count DESC
        LIMIT %s
        """
        params.append(limit)
    
        try:
            with self.conn.cursor() as cursor:
                cursor.execute(query_sql, tuple(params))
                results = cursor.fetchall()
                return results
        except Exception as e:
            print(f"查询失败: {e}")
            return []

    def insert_table_data(self, table_data_list: List[Tuple[str, dict, str, str]]):
        """插入表格数据到数据库

        Args:
            table_data_list (List[Tuple[str, dict, str, str]]): 表格数据列表，每项包含(关键词, 表格JSON数据, category_id, 源文件名)
        """
        insert_sql = """
        INSERT INTO table_storage (keywords, table_data, category_id, source_file)
        VALUES (%s, %s, %s, %s)
        ON CONFLICT DO NOTHING;
        """

        try:
            with self.conn.cursor() as cursor:
                for keywords, table_json, category_id, source_file in table_data_list:
                    cursor.execute(insert_sql, (
                        keywords,
                        Json(table_json),
                        category_id,
                        source_file
                    ))

                self.conn.commit()
                print(f"成功插入 {len(table_data_list)} 个表格数据")
                return True
        except Exception as e:
            print(f"插入表格数据失败: {e}")
            self.conn.rollback()
            raise

    def query_table_data(self, keywords: List[str], category_id: str = None, limit: int = 5):
        """根据关键词查询表格数据

        Args:
            keywords (List[str]): 关键词列表
            category_id (str, optional): 类别ID. Defaults to None.
            limit (int, optional): 返回结果数量限制. Defaults to 5.

        Returns:
            list: 查询结果
        """
        # 构建关键词匹配条件
        keyword_conditions = []
        params = []
        
        for keyword in keywords:
            keyword_conditions.append("keywords ILIKE %s")
            params.append(f"%{keyword}%")
        
        query_sql = f"""
        SELECT keywords, table_data, category_id, source_file
        FROM table_storage
        WHERE ({' OR '.join(keyword_conditions)})
        """
        
        if category_id:
            query_sql += " AND category_id = %s"
            params.append(category_id)
        
        query_sql += " ORDER BY created_at DESC LIMIT %s"
        params.append(limit)

        try:
            with self.conn.cursor() as cursor:
                cursor.execute(query_sql, tuple(params))
                results = cursor.fetchall()
                return results
        except Exception as e:
            print(f"查询表格数据失败: {e}")
            return []